#include "AnaDef.h"
#include "AnaAliDef.h"
#include "style.h"

const Double_t fthres = 200;

//move all counts to last pT bin, i.e. 10 GeV/c
void migrate(TH2D * hh)
{
  printf("migrating %s %s\n", hh->GetName(), hh->GetTitle());
  for(Int_t iy=1; iy<=hh->GetNbinsY(); iy++){
    //last at 10 GeV/c pT, so x = log10(10)
    Int_t lastx = hh->GetXaxis()->FindBin(1);
    while(hh->GetBinContent(lastx,iy)<fthres){
      printf("migrating not enough stat at lastx! %d %d %e\n", lastx, iy, hh->GetBinContent(lastx,iy));
      lastx--;
      if(hh->GetXaxis()->GetBinUpEdge(lastx)<5)
        break;
    }
    const Int_t count = hh->Integral(lastx, hh->GetNbinsX()+1, iy, iy);

    const Double_t oldn = hh->GetBinContent(lastx, iy);

    hh->SetBinContent(lastx, iy, count);
    const Double_t newn = hh->GetBinContent(lastx, iy);
    printf("migrating %s y %d x %d (%e), %e -> %e\n", hh->GetName(), iy, lastx, pow(10, hh->GetXaxis()->GetBinUpEdge(lastx)), oldn, newn);
    for(Int_t ix = lastx+1; ix<=hh->GetNbinsX()+1; ix++){
      hh->SetBinContent(ix, iy, 0);
    }
  }
}

void getHeff(const TString tag,  TH2D *hgen,  TH2D * hrec)
{
  if(tag.Contains("Pt")){
    //rebin does not affect the numbers, keep it, otherwise 11a has problem due to small stat.
    hgen->RebinX(10);
    hrec->RebinX(10);

    migrate(hgen);
    migrate(hrec);
  }

  hgen->SetName("hgen");
  hrec->SetName("hrec");
  TList *ll = new TList;
  ll->Add(hrec);
  ll->Add(hgen);

  //================== empty initialization =============
  TH2D *h2deff = (TH2D*) hrec->Clone("h2deff");
  h2deff->Scale(0);
  ll->Add(h2deff);

  TH1D *heffs[4]={0x0,0x0,0x0,0x0};
  TH1D *h2piEffs[4]={0x0,0x0,0x0,0x0};
  for(Int_t iy=1; iy<=4; iy++){
    const Int_t itype = iy-1;

    heffs[itype] = h2deff->ProjectionX(Form("eff%d", itype), iy, iy);
    heffs[itype]->Sumw2();
    heffs[itype]->Scale(0);
    ll->Add(heffs[itype]);

    h2piEffs[itype] = h2deff->ProjectionX(Form("Eff2Pi%d", itype), iy, iy);
    h2piEffs[itype]->Sumw2();
    h2piEffs[itype]->Scale(0);
    ll->Add(h2piEffs[itype]);
  }

  //======== calc ==============
  for(Int_t ix = 1; ix<=h2deff->GetNbinsX(); ix++){
    for(Int_t iy = 1; iy<=h2deff->GetNbinsY(); iy++){
      const Int_t itype = iy -1;
      const Double_t cgen = hgen->GetBinContent(ix,iy);
      if(cgen<fthres){
        continue;
      }

      const Double_t crec = hrec->GetBinContent(ix,iy);
      if(crec>cgen){
        printf("error crec>cgen ix %d iy %d crec %.0f cgen %.0f\n", ix, iy, crec, cgen);
        continue;
      }

      const Double_t eff = crec/cgen;
      const Double_t err = TMath::Sqrt(eff*(1-eff)/cgen);

      printf("check ix %d iy %d crec %.0f cgen %.0f eff %e err %e\n", ix, iy, crec, cgen, eff, err);


      h2deff->SetBinContent(ix,iy,eff);
      h2deff->SetBinError(ix,iy,err);

      heffs[itype]->SetBinContent(ix, eff);
      heffs[itype]->SetBinError(ix,err);
    }
  }
  
  for(Int_t ipar=0; ipar<4; ipar++){
    h2piEffs[ipar]->Divide(heffs[ipar], heffs[1]);
  }
  ll->ls();

  TFile *fout=new TFile(Form("OutEff-%s.root",tag.Data()), "recreate");
  ll->Write();
  fout->Save();
  fout->Close();
}

int main(int argc, char * argv[])
{
  for(int ii=0; ii<argc; ii++){
    printf("%d: %s\n", ii, argv[ii]);
  }

  if(argc!=4){
    printf("argc!=4 first histGen, then histRec, then hname: mcpdgVsPt mcpdgVsEta \n");return 1;
  }

  const TString hname(argv[3]);

  const TString tag(hname);

  TFile *fgen=new TFile(argv[1]);
  TH2D *hgen = (TH2D*)fgen->Get(hname);
  if(!hgen){
    printf("hgen null!!\n"); fgen->ls(); return 1;
  }
  /*
  TString tgen(argv[1]);
  if(tgen.Contains("/"))
    tgen=tgen(tgen.Last('/')+1, tgen.Length());
  tag+= tgen(0, tgen.Last('.'));
  */

  TFile *frec=new TFile(argv[2]);
  TH2D *hrec = (TH2D*)frec->Get(hname);
  if(!hrec){
    printf("hrec null!!\n"); frec->ls(); return 1;
  }

  /*
  TString trec(argv[2]);
  if(trec.Contains("/"))
    trec=trec(trec.Last('/')+1, trec.Length());
  tag+= trec(0, trec.Last('.'));
  printf("\ntag %s \ngen %s \nrec %s\n", tag.Data(), tgen.Data(), trec.Data());
  */
  
  getHeff(tag, hgen, hrec);

  return 1;
}

